package tokenBucket

import (
	"context"
	"fmt"
	"github.com/robfig/cron/v3"
	"go-yuque/logger"
	"time"
)

type TokenBucket struct {
	// 令牌桶容量
	Capacity int
	c        chan int
	// 令牌产生速率
	Rate time.Duration
	ctx  context.Context
	job  *cron.Cron
}

func Create(capacity int, rate time.Duration) *TokenBucket {
	t := &TokenBucket{
		Capacity: capacity,
		Rate:     rate,
		c:        make(chan int, capacity),
	}
	for range capacity {
		t.c <- 1
	}
	ctx, cancelFunc := context.WithCancel(context.Background())
	t.ctx = ctx

	go func(t *TokenBucket) {

		job := cron.New(cron.WithChain(
			cron.SkipIfStillRunning(cron.DiscardLogger),
			cron.Recover(cron.DiscardLogger),
		))
		t.job = job

		perRate := t.Rate.Seconds() / float64(t.Capacity)
		logger.Infof("tokenBucker job 创建, 每%.2f秒产生一个token", perRate)
		job.AddFunc(fmt.Sprintf("@every %fs", perRate), func() {
			select {
			case <-t.ctx.Done():
				logger.Infof("tokenBucket job stop")
				defer cancelFunc()
				job.Stop()
				return
			default:
				t.c <- 1
				logger.Infof("tokenBucket add token\n")
			}
		})

		job.Start()

	}(t)

	return t
}

func (t *TokenBucket) Close() {
	if t.job != nil {
		(*t.job).Stop()
	}
}

func (t *TokenBucket) TryTakeOne() {
	<-t.c
	logger.Debugf("tokenBucket take one token, less: %d\n", len(t.c))
}
